import os
import cv2
import numpy as np

def load_images_from_folder(folder):
    """ 加载指定文件夹下的所有 .png 图像 """
    image_files = [f for f in os.listdir(folder) if f.lower().endswith('.png')]
    image_files.sort()  # 按文件名排序

    images = []
    for filename in image_files:
        file_path = os.path.join(folder, filename)

        # 读取图像（支持中文路径）
        with open(file_path, 'rb') as f:
            img_data = np.frombuffer(f.read(), dtype=np.uint8)
            img = cv2.imdecode(img_data, cv2.IMREAD_UNCHANGED)

        images.append(img)
        print(f"已加载: {filename}")

    return images

def stitch_images(images, cols, bg_color=(0, 0, 0, 0)):
    """
    将图像按指定列数拼接成一张大图
    :param images: 图像列表
    :param cols: 每行显示的图像数量
    :param bg_color: 背景颜色（支持 RGBA）
    :return: 拼接后的图像
    """
    if not images:
        raise ValueError("没有图像可供拼接")

    h, w = images[0].shape[:2]
    rows = (len(images) + cols - 1) // cols

    # 创建空白背景图
    channels = images[0].shape[2] if len(images[0].shape) == 3 else 1
    if channels == 1:
        bg_color = (bg_color[0], )  # 灰度图
    elif channels == 3:
        bg_color = bg_color[:3]  # RGB
    elif channels == 4:
        bg_color = bg_color[:4]  # RGBA

    canvas = np.full((h * rows, w * cols, channels), bg_color, dtype=np.uint8)

    for idx, img in enumerate(images):
        row = idx // cols
        col = idx % cols
        x = col * w
        y = row * h
        canvas[y:y+h, x:x+w] = img

    return canvas

def main():
    folder = input("请输入包含PNG图像的文件夹路径：").strip()
    if not os.path.isdir(folder):
        print("输入的不是有效文件夹路径！")
        return

    try:
        cols = int(input("请输入每行显示的图像数量："))
    except ValueError:
        print("请输入有效的数字！")
        return

    output_path = input("请输入输出图像的完整路径（如 output.png）：").strip()

    images = load_images_from_folder(folder)

    if not images:
        print("未找到任何PNG图像！")
        return

    # 检查所有图像是否尺寸一致
    first_h, first_w = images[0].shape[:2]
    for img in images[1:]:
        h, w = img.shape[:2]
        if h != first_h or w != first_w:
            print("警告：图像尺寸不一致，可能影响拼接效果。")

    # 拼接图像
    print("开始拼接图像...")
    stitched_image = stitch_images(images, cols)

    # 保存结果
    cv2.imwrite(output_path, stitched_image)
    print(f"拼接完成，已保存为: {output_path}")

if __name__ == '__main__':
    main()